from .GPool import *
import requests

try:
    from httplib import HTTPConnection
    from urllib2 import HTTPHandler, build_opener
except ImportError:  # Python 3
    from http.client import HTTPConnection
    from urllib.request import HTTPHandler, build_opener, HTTPRedirectHandler


class HTTP10Connection(HTTPConnection):
    _http_vsn = 10
    _http_vsn_str = "HTTP/1.0"


# 禁止重定向
class HTTP10Handler(HTTPHandler, HTTPRedirectHandler):

    def http_error_302(self, req, fp, code, msg, headers):
        pass

    def http_open(self, req):
        return self.do_open(HTTP10Connection, req)


opener = build_opener(HTTP10Handler)


class RestfulAttack(GPool):
    def __init__(self, data, concurrent, headers=None, timeout=8):
        super().__init__(data, concurrent=concurrent)
        self.headers = headers
        self.timeout = timeout

    def get_http1_0(self, url, headers=None, timeout=None, allow_redirects=True):
        return str(opener.open(url, timeout=timeout).read(), encoding='utf-8')

    def get(self, url, headers=None, timeout=None, allow_redirects=True):
        headers = self.headers if headers is None else headers
        timeout = self.timeout if timeout is None else timeout
        try:
            r = requests.get(url=url, headers=headers, timeout=timeout, allow_redirects=allow_redirects)
            return r
        except Exception as e:
            return e

    def post(self, url, payload=None, headers=None, timeout=None, allow_redirects=True):
        headers = self.headers if headers is None else headers
        timeout = self.timeout if timeout is None else timeout
        try:
            return requests.post(url=url, json=payload, headers=headers, timeout=timeout,
                                 allow_redirects=allow_redirects)
        except Exception as e:
            return e

    def do_work(self, row):
        raise Exception("interface")

    def done(self):
        print("restful end")
